package com.soup.hummingbird.module;

import javassist.CannotCompileException;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.CtNewMethod;
import javassist.NotFoundException;

import java.io.IOException;

/**
 * 〈装载器〉<br>
 *     1:构建代理监听环境
 *     2:为目标类载入代理监听
 *
 *     基于的是javassist字节码技术
 *
 * @author zhaoyi
 * @date 2021-04-06 14:32
 * @see [相关类/方法]（可选）
 * @since [产品/模块版本] （可选）
 */
public class AgentLoader {
    private final String className;
    private final ClassLoader loader;
    private final CtClass ctClass;

    public AgentLoader(String className, ClassLoader loader, CtClass ctClass) {
        this.className = className;
        this.loader = loader;
        this.ctClass = ctClass;
    }

    /**
     * 插入 监听 method
     */
    public void updateMethod(CtMethod method, MethodSrcBuild srcBuild) throws CannotCompileException {
        CtMethod ctmethod = method;
        String methodName = method.getName();

        CtMethod agentMethod = CtNewMethod.copy(ctmethod, methodName, ctClass, null);
        agentMethod.setName(methodName + "$agent");
        ctClass.addMethod(agentMethod);

        ctmethod.setBody(srcBuild.buildSrc(ctmethod));
    }

    /**
     * 生成新的class字节码
     */
    public byte[] toBytecode() throws IOException, CannotCompileException {
        return ctClass.toBytecode();
    }

    /**
     * 插桩源码的构建器
     * 构建插桩前置源码，插桩结束源码，异常源码
     */
    public static class MethodSrcBuild {
        // 插桩前置源码
        private String beginSrc;
        // 插桩结束源码
        private String endSrc;
        // 异常源码
        private String errorSrc;

        public MethodSrcBuild setBeginSrc(String beginSrc) {
            this.beginSrc = beginSrc;
            return this;
        }

        public MethodSrcBuild setEndSrc(String endSrc) {
            this.endSrc = endSrc;
            return this;
        }

        public MethodSrcBuild setErrorSrc(String errorSrc) {
            this.errorSrc = errorSrc;
            return this;
        }

        public String buildSrc(CtMethod method) {
            String result = "";
            try {
                String template = method.getReturnType().getName().equals("void") ? voidSource : source;
                String bsrc = null == beginSrc  ? "" : beginSrc;
                String eSrc = null == errorSrc ? "" : errorSrc;
                String enSrc = null == endSrc ? "" : endSrc;
                result = String.format(template, bsrc, method.getName(), eSrc, enSrc);
            } catch (NotFoundException e) {
                throw new RuntimeException(e);
            }

            return result;
        }

        final static String source = "{\n"
                +          "%s"
                + "        Object result = null;\n"
                + "        try {\n"
                + "            result = ($w)%s$agent($$);\n"
                + "        } catch (Throwable e) {\n"
                +              "%s"
                + "            throw e;\n"
                + "        } finally {\n"
                +             "%s"
                + "        }\n"
                + "        return ($r) result;\n"
                + "}\n";

        final static String voidSource = "{\n"
                +          "%s"
                + "        try {\n"
                + "            %s$agent($$);\n"
                + "        } catch (Throwable e) {\n"
                +              "%s"
                + "            throw e;\n"
                + "        } finally {\n"
                +              "%s"
                + "        }\n"
                + "}\n";
    }

}
